open Parsetree;;
open Asttypes;;
open Longident;;

let max_width = 65536;;

type sz = No | Star | Sz of int

type t =
  | Lit of string
  | Int of string * sz * sz * char
  | Int32 of string * sz * sz * char
  | Int64 of string * sz * sz * char
  | Intnat of string * sz * sz * char
  | Float of string * sz * sz * char
  | String_s of (bool * int option) option
  | String_S of (bool * int option) option
  | Char_c
  | Char_C
  | Bool
  | Alpha
  | Beta
  | Flush
;;

let print l =
  let s oc sz = match sz with
    | No -> Printf.fprintf oc "_"
    | Star -> Printf.fprintf oc "*"
    | Sz n -> Printf.fprintf oc "%d" n
  in
  let p fmt w d c = Printf.eprintf " (%S, %a, %a, %C)" fmt s w s d c in
  let f x =
    prerr_char '<';
    begin match x with
      | Lit s -> Printf.eprintf "Lit %S" s
      | Int (fmt, w, d, c) -> Printf.eprintf "Int"; p fmt w d c
      | Int32 (fmt, w, d, c) -> Printf.eprintf "Int32"; p fmt w d c
      | Int64 (fmt, w, d, c) -> Printf.eprintf "Int64"; p fmt w d c
      | Intnat (fmt, w, d, c) -> Printf.eprintf "Intnat"; p fmt w d c
      | Float (fmt, w, d, c) -> Printf.eprintf "Float"; p fmt w d c
      | String_s None -> Printf.eprintf "String_s"
      | String_s (Some (m, Some w)) -> Printf.eprintf "String_s (%b, %d)" m w
      | String_s (Some (m, None)) -> Printf.eprintf "String_s (%b, *)" m
      | String_S None -> Printf.eprintf "String_S"
      | String_S (Some (m, Some w)) -> Printf.eprintf "String_S (%b, %d)" m w
      | String_S (Some (m, None)) -> Printf.eprintf "String_S (%b, *)" m
      | Char_c -> Printf.eprintf "Char_c"
      | Char_C -> Printf.eprintf "Char_C"
      | Bool -> Printf.eprintf "Bool"
      | Alpha -> Printf.eprintf "Alpha"
      | Beta -> Printf.eprintf "Beta"
      | Flush -> Printf.eprintf "Flush"
    end;
    prerr_string "> ";
  in
    List.iter f l;
    prerr_newline ();
;;

let parse_format fmt =
  let len = String.length fmt in
  let error i msg =
    failwith (Printf.sprintf "invalid format: %S: char %d: %s" fmt i msg)
  in
  let ueof i = error i "unexpected end of format" in
  let rec scan_format i acc =
    if i = len then
      acc
    else if fmt.[i] = '%' then
      scan_percent (succ i) acc
    else
      scan_lit i acc
  and scan_lit i acc =
    let rec f i b acc =
      if i = len then
	Lit (Buffer.contents b) :: acc
      else if fmt.[i] = '%' then
	if succ i = len then
	  ueof i
	else if fmt.[succ i] = '%' then (
	  Buffer.add_char b '%';
	  f (i + 2) b acc
	) else
	  scan_percent (succ i) (Lit (Buffer.contents b) :: acc)
      else (
	Buffer.add_char b fmt.[i];
	f (succ i) b acc
      )
    in
      f i (Buffer.create 16) acc
  and scan_percent i acc =
    let rec f i p m z acc =
      if i = len then ueof i;
      match fmt.[i] with
	| '+' -> f (succ i) true m z acc
	| '-' -> f (succ i) p true z acc
	| '0' -> f (succ i) p m true acc
	| ' ' | '#' -> f (succ i) p m z acc
	| '1' .. '9' as c ->
	    g (succ i) p m z (int_of_char c - int_of_char '0') acc
	| '.' ->
	    h (succ i) p m z No acc
	| '*' ->
	    v (succ i) p m z acc
	| _ ->
	    l i p m z No No acc
    and g i p m z w acc =
      if i = len then ueof i;
      match fmt.[i] with
	| '0' .. '9' as c ->
	    let new_w = w * 10 + int_of_char c - int_of_char '0' in
	      if new_w >= max_width then error i "too big width";
	      g (succ i) p m z new_w acc
	| '.' ->
	    h (succ i) p m z (Sz w) acc
	| _ ->
	    l i p m z (Sz w) No acc
    and v i p m z acc =
      if i = len then ueof i;
      match fmt.[i] with
	| '.' -> h (succ i) p m z Star acc
	| _ -> l i p m z Star No acc
    and h i p m z w acc =
      if i = len then ueof i;
      match fmt.[i] with
	| '0' .. '9' as c ->
	    k (succ i) p m z w (int_of_char c - int_of_char '0') acc
	| '*' ->
	    l (succ i) p m z w Star acc
	| _ ->
	    l i p m z w No acc
    and k i p m z w d acc =
      if i = len then ueof i;
      match fmt.[i] with
	| '0' .. '9' as c ->
	    let new_d = 10 * d + int_of_char c - int_of_char '0' in
	      if new_d >= max_width then error i "too big precision";
	      k (succ i) p m z w new_d acc
	| _ ->
	    l i p m z w (Sz d) acc
    and l i p m z w d acc =
      let b = Buffer.create 16 in
      let chk_void () =
	if p || m || z || w <> No || d <> No then
	  error i "unexpected arguments for this kind of format"
      in
      let extract_str () =
	if p || z || d <> No then
	  error i "unexpected arguments for string format";
	match (m, w) with
	  | (true, No) -> error i "width expected with left-justification"
	  | (false, No) -> None
	  | (m, Star) -> Some (m, None)
	  | (m, Sz w) -> Some (m, Some w)
      in
      let mk_intX b cons =
	if succ i = len then (cons (Buffer.contents b) 'd', succ i) else
	  match fmt.[succ i] with
	    | 'd' | 'i' ->
		(cons (Buffer.contents b) 'd', i + 2)
	    | 'u' | 'x' | 'X' | 'o' ->
		(cons (Buffer.contents b) fmt.[succ i], i + 2)
	    | _ ->
		(Int (Buffer.contents b, w, d, 'd'), i + 1)
      in
	if m && z && d = No then
	  error i "incompatible options: '-' and '0' with no precision";
	if m && w = No then
	  error i "width expected with left-justification";
	if z && w = No && d = No then
	  error i "incompatible options: '0' with no width nor precision";
	(match w with Sz w when w <= 1 ->
	   error i "too small width" | _ -> ());
	(match d with Sz d when d <= 1 ->
	   error i "too small precision" | _ -> ());
	begin match (w, d) with
	  | (Sz w, Sz d) when w < d -> error i "width < precision"
	  | _ -> ()
	end;
	Buffer.add_char b '%';
	if p then Buffer.add_char b '+';
	if m then Buffer.add_char b '-';
	if z then Buffer.add_char b '0';
	let (x, nxt) =
	  match fmt.[i] with
	    | 's' -> (String_s (extract_str ()), succ i)
	    | 'S' -> (String_S (extract_str ()), succ i)
	    | 'c' -> chk_void (); (Char_c, succ i)
	    | 'C' -> chk_void (); (Char_C, succ i)
	    | 'b' | 'B' -> chk_void (); (Bool, succ i)
	    | 'a' -> chk_void (); (Alpha, succ i)
	    | 't' -> chk_void (); (Beta, succ i)
	    | '!' -> chk_void (); (Flush, succ i)
	    | 'd' | 'i' | 'u' | 'x' | 'X' | 'o' ->
		(Int (Buffer.contents b, w, d, fmt.[i]), succ i)
	    | 'f' | 'F' | 'e' | 'E' | 'g' | 'G' ->
		(Float (Buffer.contents b, w, d, fmt.[i]), succ i)
	    | 'l' -> mk_intX b (fun s c -> Int32 (s, w, d, c))
	    | 'L' -> mk_intX b (fun s c -> Int64 (s, w, d, c))
	    | 'n' -> mk_intX b (fun s c -> Intnat (s, w, d, c))
	    | c -> error i (Printf.sprintf "unknown format: %%%c" c)
	in
	  scan_format nxt (x :: acc)
    in
      f i false false false acc
  in
    scan_format 0 []
;;


let extract_format fmt =
  let exp exp_desc = { pexp_desc = exp_desc; pexp_loc = Location.none } in
  let pat pat_desc = { ppat_desc = pat_desc; ppat_loc = Location.none } in
  let mk_cons cons args =
    let otpl = match args with
      | [] -> None
      | [ed] -> Some (exp ed)
      | _ -> Some (exp (Pexp_tuple (List.map exp args)))
    in
    let ident = Ldot (Ldot (Lident "Printf", "CamlinternalFormat"), cons) in
      exp (Pexp_construct (ident, otpl, false))
  in
  let mk_bool b = Pexp_construct (Lident (string_of_bool b), None, false) in
  let mk_int n = Pexp_constant (Const_int n) in
  let mk_ident i = let v = "x" ^ string_of_int i in (v, Lident v) in
  let mk_simple i params cons =
    let (v, id) = mk_ident i in
      (succ i, v :: params, mk_cons cons [ Pexp_ident id ])
  in
  let mk_bi i params cons fmt =
    let (v, id) = mk_ident i in
    let args = [ Pexp_constant (Const_string fmt) ; Pexp_ident id ] in
      (succ i, v :: params, mk_cons cons args)
  in
  let mk_quad i params cons left right =
    let (v1, id1) = mk_ident i in
    let (v2, id2) = mk_ident (succ i) in
    let args = [ Pexp_constant (Const_string left) ; Pexp_ident id2 ;
		 Pexp_constant (Const_string right) ; Pexp_ident id1 ] in
      (i + 2, v2 :: v1 :: params, mk_cons cons args)
  in
  let mk_quint i params cons left right =
    let (v1, id1) = mk_ident i in
    let (v2, id2) = mk_ident (i + 1) in
    let (v3, id3) = mk_ident (i + 2) in
    let args = [ Pexp_constant (Const_string left) ; Pexp_ident id3 ;
		 Pexp_ident id2 ; Pexp_constant (Const_string right) ;
		 Pexp_ident id1 ] in
      (i + 3, v3 :: v2 :: v1 :: params, mk_cons cons args)
  in
  let coc co c =
    match co with
      | None -> String.make 1 c
      | Some c' -> let s = String.make 2 c in s.[0] <- c'; s
  in
  let mk_gen i params cons fmt w d co c =
    match (fmt, w, d, c) with
      | ("%", No, No, ('d' | 'i' | 'f')) ->
	  mk_simple i params cons
      | (fmt, Star, No, c) ->
	  mk_quad i params (cons ^ "_sfmt") fmt (coc co c)
      | (fmt, Star, Sz d, c) ->
	  let b = Buffer.create 16 in
	    Buffer.add_char b '.';
	    Buffer.add_string b (string_of_int d);
	    Buffer.add_string b (coc co c);
	    mk_quad i params (cons ^ "_sfmt") fmt (Buffer.contents b)
      | (fmt, No, Star, c) ->
	  mk_quad i params (cons ^ "_sfmt") (fmt ^ ".") (coc co c)
      | (fmt, Sz w, Star, c) ->
	  mk_quad i params (cons ^ "_sfmt") (fmt ^ string_of_int w) (coc co c)
      | (fmt, Star, Star, c) ->
	  mk_quint i params (cons ^ "_ssfmt") fmt (coc co c)
      | (_, _, _, _) ->
	  let b = Buffer.create 16 in
	    Buffer.add_string b fmt;
	    begin match w with
	      | Sz w ->
		  Buffer.add_string b (string_of_int w)
	      | _ -> ()
	    end;
	    begin match d with
	      | Sz d ->
		  Buffer.add_char b '.';
		  Buffer.add_string b (string_of_int d)
	      | _ -> ()
	    end;
	    begin match co with
	      | Some c ->
		  Buffer.add_char b c;
	      | None -> ()
	    end;
	    Buffer.add_char b c;
	    mk_bi i params (cons ^ "_fmt") (Buffer.contents b)
  in
  let exp_of_fmt i params x = match x with
    | Lit s ->
	(i, params, mk_cons "String_s" [ Pexp_constant (Const_string s) ])
    | Int (fmt, w, d, c) ->
	mk_gen i params "Int" fmt w d None c
    | Int32 (fmt, w, d, c) ->
	mk_gen i params "Int32" fmt w d (Some 'l') c
    | Int64 (fmt, w, d, c) ->
	mk_gen i params "Int64" fmt w d (Some 'L') c
    | Intnat (fmt, w, d, c) ->
	mk_gen i params "Intnat" fmt w d (Some 'n') c
    | Float (fmt, w, d, c) ->
	mk_gen i params "Float" fmt w d None c
    | String_s None ->
	mk_simple i params "String_s"
    | String_s (Some (m, None)) ->
	let (v1, id1) = mk_ident i in
	let (v2, id2) = mk_ident (succ i) in
	let args = [ mk_bool m ; Pexp_ident id2 ; Pexp_ident id1 ] in
	  (i + 2, v2 :: v1 :: params, mk_cons "String_spad" args)
    | String_s (Some (m, Some w)) ->
	let (v, id) = mk_ident i in
	let args = [ mk_bool m ; mk_int w ; Pexp_ident id ] in
	  (succ i, v :: params, mk_cons "String_spad" args)
    | String_S None ->
	mk_simple i params "String_S"
    | String_S (Some (m, None)) ->
	let (v1, id1) = mk_ident i in
	let (v2, id2) = mk_ident (succ i) in
	let args = [ mk_bool m ; Pexp_ident id2 ; Pexp_ident id1 ] in
	  (i + 2, v2 :: v1 :: params, mk_cons "String_Spad" args)
    | String_S (Some (m, Some w)) ->
	let (v, id) = mk_ident i in
	let args = [ mk_bool m ; mk_int w ; Pexp_ident id ] in
	  (succ i, v :: params, mk_cons "String_Spad" args)
    | Char_c ->
	mk_simple i params "Char_c"
    | Char_C ->
	mk_simple i params "Char_C"
    | Bool ->
	mk_simple i params "Bool"
    | Alpha ->
	let (v1, id1) = mk_ident i in
	let (v2, id2) = mk_ident (succ i) in
	let var = "x" in
	let pat = pat (Ppat_var var) in
	let arg1 = ("", exp (Pexp_ident (Lident var))) in
	let arg2 = ("", exp (Pexp_ident id1)) in
	let body = exp (Pexp_apply (exp (Pexp_ident id2), [ arg1 ; arg2 ])) in
	let arg = Pexp_function ("", None, [ (pat, body) ]) in
	  (i + 2, v2 :: v1 :: params, mk_cons "Alpha" [ arg ])
    | Beta ->
	let (v, id) = mk_ident i in
	  (succ i, v :: params, mk_cons "Alpha" [ Pexp_ident id ])
    | Flush ->
	(i, params, mk_cons "Flush" [])
  in
  let rec fold_fmt l i params body =
    match l with
      | [] -> (params, body)
      | x :: tl ->
	  let (new_i, new_params, e) = exp_of_fmt i params x in
	  let tpl = exp (Pexp_tuple [ e ; exp body ]) in
	  let cons = (Pexp_construct (Lident "::", Some tpl, false)) in
	    fold_fmt tl new_i new_params cons
  in
  let (params, lst) =
    fold_fmt (parse_format fmt) 1 [] (Pexp_construct (Lident "[]", None, false))
  in
  let var = "k" in
  let call = Pexp_apply (exp (Pexp_ident (Lident var)), [ ("", exp lst) ]) in
  let rec mk_fun params e =
    match params with
      | [] -> e
      | param :: rest ->
	  Pexp_function ("", None, [ (pat (Ppat_var param),
				      exp (mk_fun rest e)) ])
  in
    Pexp_function ("", None, [ (pat (Ppat_var var), exp (mk_fun params call)) ])
;;
